package typechecker.walkers;

import java.util.Set;

import typechecker.Stage;
import typechecker.TypeCheckingException;
import typechecker.ast.type.ForallType;
import typechecker.ast.type.FunctionType;
import typechecker.ast.type.TypeConstructorApplication;
import typechecker.ast.type.TypeExpression;
import typechecker.ast.type.TypeVariable;

/**
 * Annotate _bare_ types with stages. Every new inference stage variable
 * used is added to the "v" set.
 */
public class TypeAnnot {

	private static void annotConstructorApplication(TypeConstructorApplication tca, Set<Integer> v, Set<Integer> vStar){
		int alpha = Stage.freshStageVar(v);
		
		if( tca.getCons().isStarred()){
			if(vStar == null){
				throw new TypeCheckingException("Starred type constructor can be used only in letrec.");
			}
			vStar.add(alpha);
		}
		
		v.add(alpha);
		tca.getCons().setStage(new Stage("v"+alpha, 0));
		
		for(TypeExpression param : tca.getParams()){
			annot(param, v, vStar);
		}
	}

	private static void annotFunctionType(FunctionType ft, Set<Integer> v, Set<Integer> vStar){
		annot(ft.getSrc(), v, vStar);
		annot(ft.getDest(), v, vStar);
	}

	private static void annotForallType(ForallType ft, Set<Integer> v, Set<Integer> vStar){
		annot(ft.getBody(), v, vStar);
	}
	
	
	private static void annotVariable(TypeVariable tv, Set<Integer> v, Set<Integer> vStar){
	}
	
	public static void annot(TypeExpression te, Set<Integer> v, Set<Integer> vStar){
		if (te instanceof ForallType){
			annotForallType((ForallType)te, v, vStar);
		} else if ( te instanceof FunctionType){
			annotFunctionType((FunctionType)te, v, vStar);
		} else if ( te instanceof TypeConstructorApplication){
			annotConstructorApplication((TypeConstructorApplication) te, v, vStar);
		} else if ( te instanceof TypeVariable){
			annotVariable((TypeVariable)te, v, vStar);
		}
	}
}
